#include <iostream>
#include <fstream>
#include <thread>
#include <chrono>

#include <glog/logging.h>

#include "cmdline.h"
#include "util.h"
#include "db.h"

using namespace std::chrono;

#define YCSB_DIR "ycsb100m"

bool opt_latency;

//==============================================================
// LOAD QUERY DATA FROM FILE
//==============================================================
void load(std::vector<_key_t> &keys, std::vector<int> &ranges, std::vector<int> &ops, std::string file) {
	std::ifstream infile_txn(file);
	if(!infile_txn) {
		fprintf(stderr, "data file is not found\n");
		exit(-1);
	}
	std::string op;
	std::string val;
	_key_t key;
	int range;
	int count = 0;
	std::string insert("INSERT");
	std::string read("READ");
	std::string update("UPDATE");
	std::string remove("REMOVE");
	std::string scan("SCAN");
	
	while (true) {
		infile_txn >> op >> key;
		if(!infile_txn.good()) {
			break;
		}
		ranges.push_back(1);
		if (op.compare(insert) == 0) {
			ops.push_back(OP_INSERT);
			keys.push_back(key);
		}
		else if (op.compare(read) == 0) {
			ops.push_back(OP_READ);
			keys.push_back(key);
		}
		else if (op.compare(update) == 0) {
			ops.push_back(OP_UPDATE);
			keys.push_back(key);
		}
		else if (op.compare(remove) == 0) {
			ops.push_back(OP_REMOVE);
			keys.push_back(key);
		}
		else if (op.compare(scan) == 0) {
			infile_txn >> range;
			ops.push_back(OP_SCAN);
			keys.push_back(key);
			ranges.back() = range;
		}
		else {
			std::cout << "UNRECOGNIZED CMD!\n";
			return;
		}
		count++;
	}

	infile_txn.close();
}

//==============================================================
// EXEC QUERIES
//==============================================================
void exec(DB * db, int num_thread, int val_size, std::vector<_key_t> &keys, 
				std::vector<int> &ranges, std::vector<int> &ops) {
	std::vector<uint64_t> latencys;
	if(opt_latency)
		latencys.resize(ops.size());

	// test part
	auto func2 = [num_thread, val_size, db, &keys, &ranges, &ops, &latencys](int thread_id) {
		size_t txn_num = ops.size();
		size_t op_per_thread = txn_num / num_thread;
		size_t start_index = op_per_thread * thread_id;
		size_t end_index = start_index + op_per_thread;

		std::string v;
		int notfound = 0;
		steady_clock::time_point last_time = steady_clock::now();
		for(size_t i = start_index;i < end_index;i++) {
			int op = ops[i];
			if (op == OP_INSERT) { //INSERT
				db->insert(keys[i], GenerateValue(keys[i], val_size));
			} else if (op == OP_READ) { //READ
				bool found = db->find(keys[i], &v);
				std::string found_key = v.substr(0, 19);
				if(found == false || keys[i] != stoll(found_key)) 
					notfound++;
			} else if (op == OP_UPDATE) { //UPDATE
				db->update(keys[i], GenerateValue(keys[i], val_size));
			} else if (op == OP_REMOVE) {
				db->remove(keys[i]);
			} else if (op == OP_SCAN) { //SCAN
				auto ret = db->scan(keys[i], ranges[i]);
				if(ranges[i] != ret) 
					notfound++;
			}

			// latency measurement
			if(opt_latency) {
				steady_clock::time_point cur_time = steady_clock::now();
				auto dur = duration_cast<nanoseconds>(cur_time - last_time);
				latencys[i] = dur.count();
				last_time = cur_time;
			}

			#ifdef DNSDEBUG
			if((i + 1) % 1000000 == 0) {
				printf("#");
				fflush(stdout);
			}
			#endif
		}
		LOG(INFO) << "not found number " << notfound;
		return;
	};

	auto start = steady_clock::now();
	
	std::thread test_threads[num_thread];
	for(int i = 0; i < num_thread; i++)
		test_threads[i] = std::move(std::thread(func2, i));

	for(int i = 0; i < num_thread; i++)
		test_threads[i].join();

	#ifdef DNSDEBUG
		printf("\n");
	#endif

	if(opt_latency) {
		std::sort(latencys.begin(), latencys.end());
		double lat50, lat90, lat99, lat999, lat9999;
		lat50 = latencys[size_t(ops.size() * 0.5)];
		lat90 = latencys[size_t(ops.size() * 0.9)];
		lat99 = latencys[size_t(ops.size() * 0.99)];
		lat999 = latencys[size_t(ops.size() * 0.999)];
		lat9999 = latencys[size_t(ops.size() * 0.9999)];

		std::cout << "min  latency:\t" << latencys[64] << std::endl
				<< "50\% latency:\t" << lat50 << std::endl
				<< "90\% latency:\t" << lat90 << std::endl
				<< "99\% latency:\t" << lat99 << std::endl
				<< "999\% latency:\t" << lat999 << std::endl
				<< "9999\% latency:\t" << lat9999 << std::endl;
	} else {
		auto end = steady_clock::now();
		auto elapse_time = duration_cast<milliseconds>(end - start);
		std::cout << (double)ops.size() / elapse_time.count() << " # Kops/s" << std::endl;
	}
	auto end_time = get_now();

	return;
}

void CleanPageCache() {
#ifndef DNSDEBUG
	auto discard = system("sync");
	discard = system("echo 3 > /proc/sys/vm/drop_caches");
	sleep(2);
	LOG(INFO) << "clean page cache";
#endif
}

int main(int argc, char *argv[]) {
	google::InitGoogleLogging(argv[0]);
	google::SetStderrLogging(google::GLOG_INFO);
	#ifndef DNSDEBUG
		google::SetLogDestination(google::GLOG_INFO, "/home/lyp/dnskv/build/INFO-");
	#endif

	cmdline::parser a;
	a.add<std::string>("dbtype", 't', "database type", true);
	a.add<int>("thread_num", 'n', "thread number", false, 1); 
	a.add<int>("value_size", 'v', "value size", false, DEFAULT_VAL_SIZE);
	a.add<std::string>("query_file", 'f', "benchmark query file", false, "/home/lyp/dnskv/build/query.dat");
	a.add<bool>("write_sync", 's', "write sync to WAL", false, false);
	a.add<bool>("latency", 'l', "measure the latency of each request", false, false);

	a.parse_check(argc, argv);

	std::string db_type = a.get<std::string>("dbtype");
	int num_thread = a.get<int>("thread_num");
	int val_size = a.get<int>("value_size");
	std::string qfile = a.get<std::string>("query_file");
	bool write_sync = a.get<bool>("write_sync");
	opt_latency = a.get<bool>("latency");

	// open or create database
	ParamType p;
	p.write_sync = write_sync;
	p.val_size = val_size;
	DB * db = getInstance(db_type, p);

	std::vector<_key_t> keys;
	std::vector<int> ranges;
	std::vector<int> ops;

	if(db_type == "matrixkv") { // test all workloads in one bench
		std::string dfile = string("/data/lyp/") + YCSB_DIR + "/dataset.dat";
		std::string qfiles[] = {
			string("/data/lyp/") + YCSB_DIR + "/query_a.dat",
			string("/data/lyp/") + YCSB_DIR + "/query_b.dat",
			string("/data/lyp/") + YCSB_DIR + "/query_c.dat",
			string("/data/lyp/") + YCSB_DIR + "/query_d.dat",
			string("/data/lyp/") + YCSB_DIR + "/query_e.dat",
			string("/data/lyp/") + YCSB_DIR + "/query_f.dat",
		};

		LOG(INFO) << "DB type:" << db_type;
		LOG(INFO) << "thread num:" << num_thread;
		LOG(INFO) << "data file:" << dfile;
		LOG(INFO) << "query file:" << qfiles[0];
		LOG(INFO) << "value size:" << val_size;
		LOG(INFO) << "latency:" << opt_latency;
		load(keys, ranges, ops, dfile);
		exec(db, 16, val_size, keys, ranges, ops);
		keys.resize(0); keys.shrink_to_fit();
		ops.resize(0); ops.shrink_to_fit();
		for(int i = 0; i < sizeof(qfiles) / sizeof(std::string); i++) {
			load(keys, ranges, ops, qfiles[i]);
			CleanPageCache();
			exec(db, num_thread, val_size, keys, ranges, ops);
			keys.resize(0); keys.shrink_to_fit();
			ops.resize(0); ops.shrink_to_fit();
		}
	} else {
		LOG(INFO) << "DB type:" << db_type;
		LOG(INFO) << "thread num:" << num_thread;
		LOG(INFO) << "query file:" << qfile;
		LOG(INFO) << "value size:" << val_size;
		LOG(INFO) << "latency:" << opt_latency;
		load(keys, ranges, ops, qfile);
		CleanPageCache();
		exec(db, num_thread, val_size, keys, ranges, ops);
	}
	
	delete db;

	google::ShutdownGoogleLogging();
	return 0;
}
